####
library(mnormt)
library(MASS)
library(tmvtnorm)

CalLik=function(beta,T,Y1,Y2,X, rho){
	 N=length(T)
	 QX=cbind(rep(1,N),T,T*X,X)
	 QW=cbind(rep(1,N),T,Y1,T*Y1,Y1*X,T*X,X)
     K=dim(QX)[2]
     L=dim(QW)[2]
     pre1=QX%*%beta[1:K]
     pre2=QW%*%beta[(K+1):(K+L)]
     Lik=0
     for (i in 1:N){
     	Lik=Lik+log(ptmvnorm(lowerx=c(ifelse(Y1[i],0,-Inf),ifelse(Y2[i],0,-Inf)),upperx=c(ifelse(Y1[i],Inf,0),ifelse(Y2[i],Inf,0)),mean =c(pre1[i],pre2[i]), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf)))
     }
     return(Lik)
}


#### EM algorithm for sensitivity analysis with rho
EM_contagion = function(T,Y1,Y2,X, rho,

                    ##parameter for EM
                      iter.max=1000,error0=0.00001,
                     
                    ##initial values
                      beta=NULL)

{
	           N=length(T)
	          QX=cbind(rep(1,N),T,T*X,X)
	          QW=cbind(rep(1,N),T,Y1,T*Y1,Y1*X,T*X,X)
              K=dim(QX)[2]
              L=dim(QW)[2]
              Y1.latent=rep(0,N)-0.5*Y1
              Y2.latent=rep(0,N)-0.5*Y2
              if(is.null(beta)) beta=rep(0,K+L)
              iter=0
              # lik=NULL
              repeat{
                     iter=iter+1
	       	  if(iter>iter.max)
	              {
	                      print("EM algorithm may not converge in the case. Please check your initial values!")
	                      converge=0
	                      break	
	              }
    
	              print(paste("The ", iter, "-th iteration!", sep=""))
	              
	              ## E step
	              pre1=QX%*%beta[1:K]
                  pre2=QW%*%beta[(K+1):(K+L)]
                 
	              for (i in 1:N){
	              	mean.Y=mtmvnorm(mean=c(pre1[i],pre2[i]),sigma=array(c(1,rho,rho,1),dim=c(2,2)),lower=c(ifelse(Y1[i],0,-Inf),ifelse(Y2[i],0,-Inf)),upper=c(ifelse(Y1[i],Inf,0),ifelse(Y2[i],Inf,0)),doComputeVariance=F)
	              	Y1.latent[i]=mean.Y$tmean[1]
	              	Y2.latent[i]=mean.Y$tmean[2]
	              }
	              ##### M step
	              sigma.inv=array(c(1,1,-1,1),dim=c(2,2))%*%diag(sqrt(c(1/(1+rho),1/(1-rho))))%*%array(c(1,-1,1,1),dim=c(2,2))/2
	              Y.trans=NULL
	              X.trans=NULL

	              for (i in 1:N){
	              	a=sigma.inv%*%c(Y1.latent[i],Y2.latent[i])
	              	Y.trans=c(Y.trans,a)
	              	b=sigma.inv%*%rbind(c(QX[i,],rep(0,L)),c(rep(0,K),QW[i,]))
	              	X.trans=rbind(X.trans,b)
	              }
	              beta.o=beta
	              beta=lm(Y.trans~0+X.trans)$coefficients
	              
	              # lik=c(lik,CalLik(beta,T,Y1,Y2,X, rho))
	              
	               error=sum((beta-beta.o)^2)
                  if(error<error0){
                      converge=1
                      break
                  }
	              }
	              return(list(beta=beta,converge=converge))
}
	
#### Calculate the effects from the parameters
CalCE=function(beta,T,Y1,Y2,X,rho){
	N=length(T)
	QX.1=cbind(rep(1,N),rep(1,N),X,X)
    QX.0=cbind(rep(1,N),rep(0,N),0*X,X)
	QW.11=cbind(rep(1,N),rep(1,N),rep(1,N),1*rep(1,N),1*X,1*X,X)
	QW.10=cbind(rep(1,N),rep(1,N),rep(0,N),1*rep(0,N),0*X,1*X,X)
	QW.01=cbind(rep(1,N),rep(0,N),rep(1,N),0*rep(1,N),1*X,0*X,X)
	QW.00=cbind(rep(1,N),rep(0,N),rep(0,N),0*rep(0,N),0*X,0*X,X)
	
	K=dim(QX.1)[2]
    L=dim(QW.00)[2]
    
	pre1.1=QX.1%*%beta[1:K]
	pre1.0=QX.0%*%beta[1:K]
	pre2.11=QW.11%*%beta[(K+1):(K+L)]
	pre2.10=QW.10%*%beta[(K+1):(K+L)]
	pre2.01=QW.01%*%beta[(K+1):(K+L)]
	pre2.00=QW.00%*%beta[(K+1):(K+L)]
	P2.11.1=rep(0,N)
	P2.11.0=rep(0,N)
	P2.10.1=rep(0,N)
	P2.10.0=rep(0,N)
	P2.01.1=rep(0,N)
	P2.01.0=rep(0,N)
	P2.00.1=rep(0,N)
	P2.00.0=rep(0,N)
	for (i in 1:N){
		P2.11.1[i]=  ptmvnorm(lowerx=c(0,0),upperx=c(Inf,Inf),mean =c(pre1.1[i],pre2.11[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.11.0[i]=  ptmvnorm(lowerx=c(0,0),upperx=c(Inf,Inf),mean =c(pre1.0[i],pre2.11[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.10.1[i]=  ptmvnorm(lowerx=c(-Inf,0),upperx=c(0,Inf),mean =c(pre1.1[i],pre2.10[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.10.0[i]=  ptmvnorm(lowerx=c(-Inf,0),upperx=c(0,Inf),mean =c(pre1.0[i],pre2.10[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		
		P2.01.1[i]= ptmvnorm(lowerx=c(0,0),upperx=c(Inf,Inf),mean =c(pre1.1[i],pre2.01[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.01.0[i]= ptmvnorm(lowerx=c(0,0),upperx=c(Inf,Inf),mean =c(pre1.0[i],pre2.01[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.00.1[i]=  ptmvnorm(lowerx=c(-Inf,0),upperx=c(0,Inf),mean =c(pre1.1[i],pre2.00[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
		P2.00.0[i]=  ptmvnorm(lowerx=c(-Inf,0),upperx=c(0,Inf),mean =c(pre1.0[i],pre2.00[i] ), sigma = array(c(1,rho,rho,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
     }
	 tau1=P2.11.1+P2.10.1-P2.11.0-P2.10.0
	 tau0=P2.01.1+P2.00.1-P2.01.0-P2.00.0
	 eta1=P2.11.1+P2.10.1-P2.01.1-P2.00.1
	 eta0=P2.11.0+P2.10.0-P2.01.0-P2.00.0
	 theta=tau1+eta0
	 
	 return(cbind(tau1,tau0,eta1,eta0,theta))
}	
	
#### Sensitivity with rho: bootstrap for pool data and Minneapolis data
BootsCE=function	(T,Y1,Y2,X,rho,K,L,start,step=10){
	 
	 N=length(T)
	 # beta.boots=array(dim=c(step,K+L))
	 # index.boots=array(dim=c(step,N))
	 ce=NULL
	 
	for (i in start:(start+step-1)){
		set.seed(i)
		ce.row=NULL
	 	print(paste("The ", i, "-th Bootstrap!", sep=""))
	 	index=sample(1:N,N,replace=TRUE)
	 	res=EM_contagion(T[index],Y1[index],Y2[index],X[index,], rho)
	 	beta.boots=res$beta
	 	ce.ind=CalCE(beta.boots,T[index],Y1[index],Y2[index],X[index,],rho=rho)
	 	ce.row=c(ce.row,apply(ce.ind,2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,1]==1&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==1&X[index,2]==0,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==0,],2,mean))
	 	ce=rbind(ce,ce.row)
	 }
	 return(ce)
}

#### EM algorithm for sensitivity analysis with sigma
EM_contagion_sens1 = function(T,Y1,Y2,X, sigma,

                    ##parameter for EM
                      iter.max=1000,error0=0.00001,
                     
                    ##initial values
                      beta=NULL)

{
	          L.X=dim(X)[2]
	           N=length(T)
	          QX=cbind(rep(1,N),T,T*X,X)
	          QW=cbind(rep(1,N),T,Y1,T*Y1,Y1*X,T*X,X)
              K=dim(QX)[2]
              L=dim(QW)[2]
              Y1.latent=rep(0,N)-0.5*Y1
              Y2.latent=rep(0,N)-0.5*Y2
              if(is.null(beta)) {beta[(K+1):(K+L)]=glm(Y2~0+QW,family=binomial(link="probit"))$coefficients}
              
              beta[1:K]=glm(Y1~0+QX,family=binomial(link="probit"))$coefficients*sqrt(1+sigma^2)
              
              pre1=rep(0,N)
              ###    p1 P(Y1*=1| T X); pY1Ys1  P(Y1, Y1*=1| T X); pY1Ys0 P(Y1, Y1*=0| T X)
              
              pY1Ys1=rep(0,N)
              pY1Ys0=rep(0,N)
              for (i in 1:N){
              	pre1[i]=sum(QX[i,]*beta[1:K])
              	
              	if (Y1[i]==1){
              		pY1Ys1[i]=ptmvnorm(lowerx=c(0,0),upperx=c(Inf,Inf),mean =c(pre1[i],pre1[i] ), sigma = array(c(1+sigma^2,1,1,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
              		pY1Ys0[i]=ptmvnorm(lowerx=c(0,-Inf),upperx=c(Inf,0),mean =c(pre1[i],pre1[i] ), sigma = array(c(1+sigma^2,1,1,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
              	}else{
              		pY1Ys1[i]=ptmvnorm(lowerx=c(-Inf,0),upperx=c(0,Inf),mean =c(pre1[i],pre1[i] ), sigma = array(c(1+sigma^2,1,1,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
              		pY1Ys0[i]=ptmvnorm(lowerx=c(-Inf,-Inf),upperx=c(0,0),mean =c(pre1[i],pre1[i] ), sigma = array(c(1+sigma^2,1,1,1),dim=c(2,2)),lower = c(-Inf, -Inf), upper = c(Inf, Inf))
              		}
              	}	
              		
              iter=0
              # lik=NULL
             
              repeat{
                     iter=iter+1
	       	  if(iter>iter.max)
	              {
	                      print("EM algorithm may not converge in the case. Please check your initial values!")
	                      converge=0
	                      break	
	              }
    
	              print(paste("The ", iter, "-th iteration!", sep=""))
	              
	              ## E step
	              mat=array(0,dim=c(L,L))
                  vec=rep(0,L)
	              rec=rep(0,N)
	               for (i in 1:N){
	              	pre2.1=sum(c(1,T[i],1,T[i]*1,1*X[i,],T[i]*X[i,],X[i,])*beta[(K+1):(K+L)])
	                pre2.0=sum(c(1,T[i],0,T[i]*0,0*X[i,],T[i]*X[i,],X[i,])*beta[(K+1):(K+L)])
	                
	                ### p  P(Y1*=1 Y1 Y2 | T X)  q  P(Y1*=0 Y1 Y2 | T X)
	                if (Y2[i]==1){
	               p=pY1Ys1[i]*pnorm(pre2.1)
	               q=pY1Ys0[i]*pnorm(pre2.0)
	                }else{
		           p=pY1Ys1[i]*pnorm(-pre2.1)
	               q=pY1Ys0[i]*pnorm(-pre2.0)
		            }
		            ###  ey1s  P(Y1*=1 | obs); ey2oy1s E(Y2^o Y1* | obs); ey2o E(Y2^o| obs)
		            ey1s=p/(p+q)
		            if (ey1s==0){
	                 ey2oy1s=0
	                 ey2o=ey2oy1s+(1-ey1s)*mtmvnorm(pre2.0,sigma=1,lower=ifelse(Y2[i],0,-Inf),upper=ifelse(Y2[i],Inf,0),doComputeVariance=FALSE)$tmean
	}
	if(ey1s==1){
		ey2oy1s=ey1s*mtmvnorm(pre2.1,sigma=1,lower=ifelse(Y2[i],0,-Inf),upper=ifelse(Y2[i],Inf,0),doComputeVariance=FALSE)$tmean
	ey2o=ey2oy1s
	}
	if(ey1s!=0&ey1s!=1){
		ey2oy1s=ey1s*mtmvnorm(pre2.1,sigma=1,lower=ifelse(Y2[i],0,-Inf),upper=ifelse(Y2[i],Inf,0),doComputeVariance=FALSE)$tmean
	ey2o=ey2oy1s+(1-ey1s)*mtmvnorm(pre2.0,sigma=1,lower=ifelse(Y2[i],0,-Inf),upper=ifelse(Y2[i],Inf,0),doComputeVariance=FALSE)$tmean
	}
		            mat.ind=array(dim=c(L,L))
		            
	mat.ind[1:2,]=c(1,T[i])%*%t(c(1,T[i],ey1s,T[i]*ey1s,ey1s*X[i,],T[i]*X[i,],X[i,]))
	mat.ind[3:(4+L.X),]=c(ey1s,T[i]*ey1s,ey1s*X[i,])%*%t(c(1,T[i],1,T[i]*1,1*X[i,],T[i]*X[i,],X[i,]))
	mat.ind[(5+L.X):L,]=c(T[i]*X[i,],X[i,])%*%t(c(1,T[i],ey1s,T[i]*ey1s,ey1s*X[i,],T[i]*X[i,],X[i,]))

	vec.ind=c(ey2o,T[i]*ey2o,ey2oy1s,T[i]*ey2oy1s,ey2oy1s*X[i,],T[i]*X[i,]*ey2o,X[i,]*ey2o)
	              	mat=mat+mat.ind
	              	vec=vec+vec.ind
		            
	                }
	              ##### M step
	              beta.o=beta
	              beta[(K+1):(K+L)]=solve(mat,vec)
	              
	              
	              # lik=c(lik,CalLik(beta,T,Y1,Y2,X, rho))
	              
	               error=sum((beta-beta.o)^2)/(K+L)^2
	               # print(error)
                  if(error<error0){
                      converge=1
                      break
                  }
	              }
	              return(list(beta=beta,converge=converge))
}


#### Sensitivity with sigma: bootstrap for pool data and Minneapolis data
BootsCE_sens1_pool=function	(T,Y1,Y2,X,sigma,K,L,start,step=10){
	 
	 N=length(T)
	 # beta.boots=array(dim=c(step,K+L))
	 # index.boots=array(dim=c(step,N))
	 ce=NULL
	 
	for (i in start:(start+step-1)){
		set.seed(i)
		ce.row=NULL
	 	print(paste("The ", i, "-th Bootstrap!", sep=""))
	 	index=sample(1:N,N,replace=TRUE)
	 	res=EM_contagion_sens1(T[index],Y1[index],Y2[index],X[index,], sigma=sigma)
	 	beta.boots=res$beta
	 	ce.ind=CalCE(beta.boots,T[index],Y1[index],Y2[index],X[index,],rho=0)
	 	ce.row=c(ce.row,apply(ce.ind,2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,1]==1&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==1&X[index,2]==0,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==0,],2,mean))
	 	ce=rbind(ce,ce.row)
	 }
	 return(ce)
}

######## Sensitivity with rho: bootstrap for Denver data
BootsCE_denver=function	(T,Y1,Y2,X,rho,K,L,start,step=10){
	 
	 N=length(T)
	 # beta.boots=array(dim=c(step,K+L))
	 # index.boots=array(dim=c(step,N))
	 ce=NULL
	 
	for (i in start:(start+step-1)){
		set.seed(i)
		ce.row=NULL
	 	print(paste("The ", i, "-th Bootstrap!", sep=""))
	 	index=sample(1:N,N,replace=TRUE)
	 	res=EM_contagion(T[index],Y1[index],Y2[index],X[index,], rho)
	 	beta.boots=res$beta
	 	ce.ind=CalCE(beta.boots,T[index],Y1[index],Y2[index],X[index,],rho=rho)
	 	ce.row=c(ce.row,apply(ce.ind,2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,1]==1&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==1&X[index,2]==0,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==0,],2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,3]==1&X[index,4]==1,],2,mean),apply(ce.ind[X[index,3]==1&X[index,4]==0,],2,mean),apply(ce.ind[X[index,3]==0&X[index,4]==1,],2,mean),apply(ce.ind[X[index,3]==0&X[index,4]==0,],2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,5]==1&X[index,6]==1,],2,mean),apply(ce.ind[X[index,5]==1&X[index,6]==0,],2,mean),apply(ce.ind[X[index,5]==0&X[index,6]==1,],2,mean),apply(ce.ind[X[index,5]==0&X[index,6]==0,],2,mean))
	 	ce=rbind(ce,ce.row)
	 }
	 return(ce)
}



######## Sensitivity with sigma: bootstrap for Denver data
BootsCE_sens1_denver=function	(T,Y1,Y2,X,sigma,K,L,start,step=10){
	 
	 N=length(T)
	 # beta.boots=array(dim=c(step,K+L))
	 # index.boots=array(dim=c(step,N))
	 ce=NULL
	 
	for (i in start:(start+step-1)){
		set.seed(i)
		ce.row=NULL
	 	print(paste("The ", i, "-th Bootstrap!", sep=""))
	 	index=sample(1:N,N,replace=TRUE)
	 	res=EM_contagion_sens1(T[index],Y1[index],Y2[index],X[index,], sigma=sigma)
	 	beta.boots=res$beta
	 	ce.ind=CalCE(beta.boots,T[index],Y1[index],Y2[index],X[index,],rho=0)
	 	ce.row=c(ce.row,apply(ce.ind,2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,1]==1&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==1&X[index,2]==0,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==1,],2,mean),apply(ce.ind[X[index,1]==0&X[index,2]==0,],2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,3]==1&X[index,4]==1,],2,mean),apply(ce.ind[X[index,3]==1&X[index,4]==0,],2,mean),apply(ce.ind[X[index,3]==0&X[index,4]==1,],2,mean),apply(ce.ind[X[index,3]==0&X[index,4]==0,],2,mean))
	 	ce.row=c(ce.row,apply(ce.ind[X[index,5]==1&X[index,6]==1,],2,mean),apply(ce.ind[X[index,5]==1&X[index,6]==0,],2,mean),apply(ce.ind[X[index,5]==0&X[index,6]==1,],2,mean),apply(ce.ind[X[index,5]==0&X[index,6]==0,],2,mean))
	 	ce=rbind(ce,ce.row)
	 }
	 return(ce)
}
	    
  

